package ssh

import (
	"encoding/json"
	"github.com/Unixeno/TheMoon/logging"
	"github.com/Unixeno/TheMoon/model"
	"github.com/Unixeno/TheMoon/moon/core/context"
	"github.com/Unixeno/TheMoon/moon/core/module"
	"github.com/Unixeno/TheMoon/pool"
	"github.com/creack/pty"
	"os"
	"os/exec"
)

const ModuleName = "ssh"

var logger = logging.GetLogger(ModuleName)

type session struct {
}

func GetModule() module.Module {
	return &session{}
}

func (s *session) Init() {

}

func (s *session) Start() {
	context.RegisterStream(ModuleName, 0, s.session)
}

type winSize struct {
	Cols uint16 `json:"cols"`
	Rows uint16 `json:"rows"`
}

func (s *session) session(stream model.Stream) error {
	cmd := exec.Command("bash")
	ptmx, err := pty.Start(cmd)

	if err != nil {
		logger.WithError(err).Error()
		return nil
	}
	defer ptmx.Close()
	pty.InheritSize(os.Stdin, ptmx)

	go func() {
		for {
			buf, err := stream.Recv()
			if err != nil {
				logger.WithError(err).Error()
				return
			}
			if len(buf) > 1 {
				if buf[0] == 'c' {
					// 命令
					size := winSize{}
					if err = json.Unmarshal(buf[1:], &size); err == nil {
						logger.WithField("col", size.Cols).WithField("rows", size.Rows).Info("reset size")
						err = pty.Setsize(ptmx, &pty.Winsize{
							Rows: size.Rows,
							Cols: size.Cols,
						})
						if err != nil {
							logger.WithError(err).Error("resize tty failed")
						}
					} else {
						logger.WithField("cmd", string(buf)).WithError(err).Error("invalid command")
					}
				} else {
					ptmx.Write(buf[1:])
				}
			}
		}
	}()
	for {
		newBuf := pool.Fixed64KPool.Get().([]byte)
		n, err := ptmx.Read(newBuf)
		if err != nil {
			logger.WithError(err).Error()
			return nil
		}
		stream.Send(newBuf[:n])
		pool.Fixed64KPool.Put(newBuf)
	}
	return nil
}
